// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <cmath>

#include <nebula/compute/api_scalar.h>
#include <nebula/compute/kernels/common_internal.h>

#include <nebula/bits/bit_util.h>
#include <nebula/bits/bitmap_ops.h>

namespace nebula::compute::internal {
    namespace {

        turbo::Status IsValidExec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
            const ArraySpan &arr = batch[0].array;
            ArraySpan *out_span = out->array_span_mutable();
            if (arr.type->id() == Type::NA) {
                // Input is all nulls => output is entirely false.
                bit_util::SetBitsTo(out_span->buffers[1].data, out_span->offset, out_span->length,
                                    false);
                return turbo::OkStatus();
            }

            DKCHECK_EQ(out_span->offset, 0);
            DKCHECK_LE(out_span->length, arr.length);
            if (arr.may_have_nulls()) {
                // We could do a zero-copy optimization, but it isn't worth the added complexity
                ::nebula::internal::CopyBitmap(arr.buffers[0].data, arr.offset, arr.length,
                                               out_span->buffers[1].data, out_span->offset);
            } else {
                // Input has no nulls => output is entirely true.
                bit_util::SetBitsTo(out_span->buffers[1].data, out_span->offset, out_span->length,
                                    true);
            }
            return turbo::OkStatus();
        }

        struct IsFiniteOperator {
            template<typename OutType, typename InType>
            static constexpr OutType Call(KernelContext *, const InType &value, turbo::Status *) {
                return std::isfinite(value);
            }
        };

        struct IsInfOperator {
            template<typename OutType, typename InType>
            static constexpr OutType Call(KernelContext *, const InType &value, turbo::Status *) {
                return std::isinf(value);
            }
        };

        using NanOptionsState = OptionsWrapper<NullOptions>;

        template<typename T>
        static void SetNanBits(const ArraySpan &arr, uint8_t *out_bitmap, int64_t out_offset) {
            const T *data = arr.get_values<T>(1);
            for (int64_t i = 0; i < arr.length; ++i) {
                if (std::isnan(data[i])) {
                    bit_util::SetBit(out_bitmap, i + out_offset);
                }
            }
        }

        turbo::Status IsNullExec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
            const ArraySpan &arr = batch[0].array;
            ArraySpan *out_span = out->array_span_mutable();
            if (arr.type->id() == Type::NA) {
                bit_util::SetBitsTo(out_span->buffers[1].data, out_span->offset, out_span->length,
                                    true);
                return turbo::OkStatus();
            }

            const auto &options = NanOptionsState::Get(ctx);
            uint8_t *out_bitmap = out_span->buffers[1].data;
            if (arr.get_null_count() > 0) {
                // Input has nulls => output is the inverted null (validity) bitmap.
                nebula::internal::InvertBitmap(arr.buffers[0].data, arr.offset, arr.length, out_bitmap,
                             out_span->offset);
            } else {
                // Input has no nulls => output is entirely false.
                bit_util::SetBitsTo(out_bitmap, out_span->offset, out_span->length, false);
            }

            if (is_floating(arr.type->id()) && options.nan_is_null) {
                switch (arr.type->id()) {
                    case Type::FP32:
                        SetNanBits<float>(arr, out_bitmap, out_span->offset);
                        break;
                    case Type::FP64:
                        SetNanBits<double>(arr, out_bitmap, out_span->offset);
                        break;
                    default:
                        return turbo::unimplemented_error("NaN detection not implemented for type ",
                                                          arr.type->to_string());
                }
            }
            return turbo::OkStatus();
        }

        struct IsNanOperator {
            template<typename OutType, typename InType>
            static constexpr OutType Call(KernelContext *, const InType &value, turbo::Status *) {
                return std::isnan(value);
            }
        };

        void MakeFunction(std::string name, FunctionDoc doc, std::vector<InputType> in_types,
                          OutputType out_type, ArrayKernelExec exec, FunctionRegistry *registry,
                          NullHandling::type null_handling, bool can_write_into_slices,
                          const FunctionOptions *default_options = nullptr,
                          KernelInit init = nullptr) {
            Arity arity{static_cast<int>(in_types.size())};
            auto func =
                    std::make_shared<ScalarFunction>(name, arity, std::move(doc), default_options);

            ScalarKernel kernel(std::move(in_types), out_type, exec, init);
            kernel.null_handling = null_handling;
            kernel.can_write_into_slices = can_write_into_slices;

            KCHECK_OK(func->add_kernel(std::move(kernel)));
            KCHECK_OK(registry->add_function(std::move(func)));
        }

        template<typename InType, typename Op>
        void AddFloatValidityKernel(const std::shared_ptr<DataType> &ty, ScalarFunction *func) {
            KCHECK_OK(func->add_kernel({ty}, boolean(),
                                      applicator::ScalarUnary<BooleanType, InType, Op>::Exec));
        }

        template<bool kConstant>
        turbo::Status ConstBoolExec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
            ArraySpan *array = out->array_span_mutable();
            bit_util::SetBitsTo(array->buffers[1].data, array->offset, array->length, kConstant);
            return turbo::OkStatus();
        }

        std::shared_ptr<ScalarFunction> MakeIsFiniteFunction(std::string name, FunctionDoc doc) {
            auto func = std::make_shared<ScalarFunction>(name, Arity::Unary(), std::move(doc));

            AddFloatValidityKernel<Fp32Type, IsFiniteOperator>(float32(), func.get());
            AddFloatValidityKernel<Fp64Type, IsFiniteOperator>(float64(), func.get());

            for (const auto &ty: int_types()) {
                KCHECK_OK(func->add_kernel({InputType(ty->id())}, boolean(), ConstBoolExec<true>));
            }
            KCHECK_OK(func->add_kernel({InputType(Type::NA)}, boolean(), ConstBoolExec<true>));
            KCHECK_OK(
                    func->add_kernel({InputType(Type::DECIMAL128)}, boolean(), ConstBoolExec<true>));
            KCHECK_OK(
                    func->add_kernel({InputType(Type::DECIMAL256)}, boolean(), ConstBoolExec<true>));
            KCHECK_OK(func->add_kernel({InputType(Type::DURATION)}, boolean(), ConstBoolExec<true>));

            return func;
        }

        std::shared_ptr<ScalarFunction> MakeIsInfFunction(std::string name, FunctionDoc doc) {
            auto func = std::make_shared<ScalarFunction>(name, Arity::Unary(), std::move(doc));

            AddFloatValidityKernel<Fp32Type, IsInfOperator>(float32(), func.get());
            AddFloatValidityKernel<Fp64Type, IsInfOperator>(float64(), func.get());

            for (const auto &ty: int_types()) {
                KCHECK_OK(func->add_kernel({InputType(ty->id())}, boolean(), ConstBoolExec<false>));
            }
            KCHECK_OK(func->add_kernel({InputType(Type::NA)}, boolean(), ConstBoolExec<false>));
            KCHECK_OK(
                    func->add_kernel({InputType(Type::DECIMAL128)}, boolean(), ConstBoolExec<false>));
            KCHECK_OK(
                    func->add_kernel({InputType(Type::DECIMAL256)}, boolean(), ConstBoolExec<false>));
            KCHECK_OK(
                    func->add_kernel({InputType(Type::DURATION)}, boolean(), ConstBoolExec<false>));
            return func;
        }

        std::shared_ptr<ScalarFunction> MakeIsNanFunction(std::string name, FunctionDoc doc) {
            auto func = std::make_shared<ScalarFunction>(name, Arity::Unary(), std::move(doc));

            AddFloatValidityKernel<Fp32Type, IsNanOperator>(float32(), func.get());
            AddFloatValidityKernel<Fp64Type, IsNanOperator>(float64(), func.get());

            for (const auto &ty: int_types()) {
                KCHECK_OK(func->add_kernel({InputType(ty->id())}, boolean(), ConstBoolExec<false>));
            }
            KCHECK_OK(func->add_kernel({InputType(Type::NA)}, boolean(), ConstBoolExec<false>));
            KCHECK_OK(
                    func->add_kernel({InputType(Type::DECIMAL128)}, boolean(), ConstBoolExec<false>));
            KCHECK_OK(
                    func->add_kernel({InputType(Type::DECIMAL256)}, boolean(), ConstBoolExec<false>));
            KCHECK_OK(
                    func->add_kernel({InputType(Type::DURATION)}, boolean(), ConstBoolExec<false>));

            return func;
        }

        turbo::Status TrueUnlessNullExec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
            ArraySpan *out_span = out->array_span_mutable();
            if (out_span->buffers[0].data) {
                // If there is a validity bitmap computed above the kernel
                // invocation, we copy it to the output buffers
                ::nebula::internal::CopyBitmap(out_span->buffers[0].data, out_span->offset,
                                               out_span->length, out_span->buffers[1].data,
                                               out_span->offset);
            } else {
                // But for all-valid inputs, the engine will skip allocating a
                // validity bitmap, so we set everything to true
                bit_util::SetBitsTo(out_span->buffers[1].data, out_span->offset, out_span->length,
                                    true);
            }
            return turbo::OkStatus();
        }

        const FunctionDoc is_valid_doc(
                "Return true if non-null",
                ("For each input value, emit true iff the value is valid (i.e. non-null)."),
                {"values"});

        const FunctionDoc is_finite_doc(
                "Return true if value is finite",
                ("For each input value, emit true iff the value is finite\n"
                 "(i.e. neither NaN, inf, nor -inf)."),
                {"values"});

        const FunctionDoc is_inf_doc(
                "Return true if infinity",
                ("For each input value, emit true iff the value is infinite (inf or -inf)."),
                {"values"});

        const FunctionDoc is_null_doc(
                "Return true if null (and optionally NaN)",
                ("For each input value, emit true iff the value is null.\n"
                 "True may also be emitted for NaN values by setting the `nan_is_null` flag."),
                {"values"}, "NullOptions");

        const FunctionDoc true_unless_null_doc("Return true if non-null, else return null",
                                               ("For each input value, emit true iff the value\n"
                                                "is valid (non-null), otherwise emit null."),
                                               {"values"});

        const FunctionDoc is_nan_doc("Return true if NaN",
                                     ("For each input value, emit true iff the value is NaN."),
                                     {"values"});

    }  // namespace

    void register_scalar_validity(FunctionRegistry *registry) {
        static auto kNullOptions = NullOptions::defaults();
        MakeFunction("is_valid", is_valid_doc, {InputType::Any()}, boolean(), IsValidExec,
                     registry, NullHandling::OUTPUT_NOT_NULL,
                /*can_write_into_slices=*/false);

        MakeFunction("is_null", is_null_doc, {InputType::Any()}, boolean(), IsNullExec,
                     registry, NullHandling::OUTPUT_NOT_NULL,
                /*can_write_into_slices=*/true, &kNullOptions, NanOptionsState::init);

        MakeFunction("true_unless_null", true_unless_null_doc, {InputType::Any()}, boolean(),
                     TrueUnlessNullExec, registry, NullHandling::INTERSECTION,
                /*can_write_into_slices=*/false);

        KCHECK_OK(registry->add_function(MakeIsFiniteFunction("is_finite", is_finite_doc)));
        KCHECK_OK(registry->add_function(MakeIsInfFunction("is_inf", is_inf_doc)));
        KCHECK_OK(registry->add_function(MakeIsNanFunction("is_nan", is_nan_doc)));
    }

}  // namespace nebula::compute::internal
